#!/usr/bin/env python
# coding:utf-8
# created by Peixuan Shu
# 2023.4.7
"""
Functions:
1. listen to uav states and check if it is connected
2. listen to /command, /land_num and /takeoff_num to takeoff or land
3. read from mocap/uwb source and forward the pose data to /telloi/pose
"""

# import numpy as np
# from math import *
# from enum import Enum
import time
# import os
from threading import Thread
import subprocess
import yaml

import rospy
from std_msgs.msg import Int32
from geometry_msgs.msg import PoseStamped
from sensor_msgs.msg import Imu
from geometry_msgs.msg import Quaternion
from std_msgs.msg import Empty
from tello_driver.msg import TelloStatus

# from tf.transformations import euler_from_quaternion

class UAVnode():
    """
    UAV communication node
    """
    def __init__(self, id, config_path, main_gui):
        self.id = id
        self.main_gui = main_gui
        self.print_console = self.main_gui.print_console
        self.warning_msg = self.main_gui.warning_msg
        self.tello_yaml = config_path
        ############ states initialize ############
        self.offline_sec = 0.8
        self.lose_pose_sec = 0.3
        self.offline_flag = True
        self.pose_source_ok = False
        self.battery_perct = -999.0  
        self.height_tello_sensor = 0.0
        self.imu_quat = Quaternion()
        self.pose = PoseStamped()
        self.last_battery_time = 0.0
        self.last_imu_time = 0.0
        self.last_pose_time = 0.0
        ############ parse tello_config ###########
        self.ip = "not found"
        self.pose_source = "not found"
        with open(self.tello_yaml, 'r') as ymlfile:
            cfg = yaml.safe_load(ymlfile)
        for uav in cfg["uavs"]:
            uav_id = int(uav["id"])
            if uav_id == self.id: ## parse uav info
                self.ip = uav["ip"]
                self.pose_source = uav["pose_source"]
                if self.pose_source == "mocap":
                    self.mocap_markerset = uav["mocap_markerset"]
                elif self.pose_source == "nluwb":
                    self.nluwb_tag_id = uav["nluwb_tag_id"]
                elif self.pose_source == "none":
                    pass
                else:
                    self.print_console("[Error] invalid pose_source: {} for tello {}".format(self.pose_source, self.id))
                break  
        if self.ip == "not found":
            # self.main_gui.topframe.uav_button[self.id].configure(bg="red") 
            self.warning_msg("[Error] id {} not found in config file: {}".format(self.id, self.tello_yaml))
        ############ ROS initialize ############
        try:
            rospy.init_node('tellogc_node', anonymous=True)
        except:
            pass # avoid overinit
        # self.print_console("加载uavnode{}".format(self.id))
        prefix = "tello"+str(self.id)
        ############## ROS topics ##############
        self.subs = []; self.pubs = []
        self.subs.append(rospy.Subscriber('/command', Int32, self.command_callback, queue_size=1))
        self.subs.append(rospy.Subscriber('/land_num', Int32, self.land_num_callback, queue_size=1))
        self.subs.append(rospy.Subscriber('/takeoff_num', Int32, self.takeoff_num_callback, queue_size=1))
        self.subs.append(rospy.Subscriber(prefix+'/status', TelloStatus, self.status_callback, queue_size=1))
        self.subs.append(rospy.Subscriber(prefix+'/imu', Imu, self.imu_callback, queue_size=1))
        if self.pose_source == "mocap":
            self.subs.append(rospy.Subscriber("/vrpn_client_node/" + self.mocap_markerset + "/pose", PoseStamped, self.mocap_callback, queue_size=1))
        elif self.pose_source == "nluwb":
            import nlink_parser.msg ## rely on nooploop uwb ros driver package
            self.subs.append(rospy.Subscriber("/nlink_linktrack_anchorframe0", nlink_parser.msg.LinktrackAnchorframe0, self.nluwb_callback, queue_size=1))
        self.pose_pub = rospy.Publisher(prefix+'/pose', PoseStamped, queue_size=10)
        self.takeoff_pub = rospy.Publisher(prefix+'/takeoff', Empty, queue_size=10)
        self.land_pub = rospy.Publisher(prefix+'/land', Empty, queue_size=10)
        self.pubs.append(self.pose_pub)
        self.pubs.append(self.takeoff_pub)
        self.pubs.append(self.land_pub)
        ############ status check thread #########
        self.kill_thread = False
        self.check_thread = Thread(target=self.check_status)
        self.check_thread.start()

    def __del__(self):
        for subcriber in self.subs:
            subcriber.unregister()
        for publisher in self.pubs:
            publisher.unregister()
        self.kill_thread = True
        del self

    def check_status(self):
        count = 0
        while not self.kill_thread:
            """ check if uav is connected by battery percentage data and imu data streaming """
            if (time.time() - self.last_battery_time > self.offline_sec) or (time.time() - self.last_imu_time > self.offline_sec):
                if self.offline_flag == False:
                    self.main_gui.topframe.uav_button[self.id].configure(bg="yellow") 
                    self.print_console("uav{} is offline!".format(self.id))   
                self.offline_flag = True     
            else:
                if self.offline_flag == True:
                    self.main_gui.topframe.uav_button[self.id].configure(bg="green") 
                    self.print_console("uav{} is online now!".format(self.id))     
                self.offline_flag = False  
            """ check if pose_source is streaming pose data (mocap or uwb) """
            if time.time() - self.last_pose_time > self.lose_pose_sec:
                self.pose_source_ok = False
            else:
                self.pose_source_ok = True
            """ check whether tello_driver_node is timeout and exit """   
            if 'tello_driver_node'+str(self.id) not in subprocess.getoutput("rosnode list"):
                count = count + 1
            if count >= 10: # tello_driver_node is timeout for 1 s
                defaultbg = self.main_gui.root.cget('bg') # default button color
                self.main_gui.topframe.uav_button[self.id].configure(bg=defaultbg)
                self.print_console("tello_driver_node{} is timeout!".format(self.id))
                # kill the tello_driver_node
                if self.id in self.main_gui.tellodriver_nodes:
                    self.main_gui.tellodriver_nodes[self.id].terminate()
                self.__del__() # delete this object
            time.sleep(0.1)
                    

    def status_callback(self,msg):
        self.battery_perct = msg.battery_percentage
        self.height_tello_sensor = msg.height_m # height from tello onboard sensor (not precise)
        self.main_gui.topframe.show_status(self.id, self.battery_perct, self.pose_source_ok, self.pose_source)    
        self.last_battery_time =time.time()

    def imu_callback(self,msg):
        """ orientation quaternion from tello onboard sensor"""
        self.imu_quat = msg.orientation # yaw is relative to the turn-on moment heading
        self.last_imu_time =time.time()

    def command_callback(self,msg):
        command = msg.data
        if command == 1:
            for i in range(10): # loop for 10 times
                pub_msg = Empty()
                self.takeoff_pub.publish(pub_msg)
                # time.sleep(0.1)
        elif command == 9:
            for i in range(20): # loop for 20 times
                pub_msg = Empty()
                self.land_pub.publish(pub_msg)
                # time.sleep(0.1)        

    def land_num_callback(self,msg):
        num = msg.data
        if num == self.id:
            for i in range(20): # loop for 20 times
                pub_msg = Empty()
                self.land_pub.publish(pub_msg)
                # time.sleep(0.1)  

    def takeoff_num_callback(self,msg):
        num = msg.data
        if num == self.id:
            for i in range(10): # loop for 10 times
                pub_msg = Empty()
                self.takeoff_pub.publish(pub_msg)
                # time.sleep(0.1)     
    
    def mocap_callback(self,data):
        pub_msg = data # foward vrpn data to pose publisher
        pub_msg.header.stamp = rospy.Time.now()
        self.pose_pub.publish(pub_msg)
        # self.state_pos[0] = data.pose.position.x
        # self.state_pos[1] = data.pose.position.y
        # self.state_pos[2] = data.pose.position.z
        # qw = data.pose.orientation.w
        # qx = data.pose.orientation.x
        # qy = data.pose.orientation.y 
        # qz = data.pose.orientation.z 
        # euler = euler_from_quaternion([qx, qy, qz, qw])
        # self.yaw = euler[2]        
        self.last_pose_time =time.time()

    def nluwb_callback(self,msg):
        nodes = msg.nodes
        for node in nodes:
            tag_id = node.id
            if tag_id == self.nluwb_tag_id:
                pub_msg = PoseStamped()
                pub_msg.header.stamp = rospy.Time.now()
                pub_msg.pose.position.x = node.pos_3d[0] # uwb 
                pub_msg.pose.position.y = node.pos_3d[1] # uwb
                pub_msg.pose.position.z = self.height_tello_sensor # height from tello onboard sensor (not precise)
                pub_msg.pose.orientation = self.imu_quat # orientation quaternion from tello onboard sensor
                self.pose_pub.publish(pub_msg)
                self.last_pose_time =time.time()
                break